import yaml
from typing import Dict, Any, List
import os
import torch
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import socket
import json
import time

class PredictionParameters:
    """
    预测参数管理类，支持从YAML文件读取参数并自动填充到界面控件
    """

    def __init__(self):
        # 默认预测参数
        self.params = {
            'project': 'runs/predict',
            'name': 'exp',
            'show': True,
            'save_frames': False,
            'save_txt': False,
            'save_conf': False,
            'save_crop': False,
            'show_labels': True,
            'show_conf': True,
            'show_boxes': True
        }

        # 存储UI控件引用（在实际应用中需要设置）
        self.ui_controls = {}

    def load_from_yaml(self, file_path: str) -> bool:
        """
        从YAML文件加载预测参数

        Args:
            file_path: YAML文件路径

        Returns:
            bool: 加载成功返回True，失败返回False
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                print(f"错误: 文件不存在 - {file_path}")
                return False

            # 读取YAML文件
            with open(file_path, 'r', encoding='utf-8') as file:
                config = yaml.safe_load(file)

            # 更新参数字典
            if config and isinstance(config, dict):
                for key, value in config.items():
                    if key in self.params:
                        self.params[key] = value
                        print(f"加载参数: {key} = {value}")
                    else:
                        print(f"警告: 未知参数 {key}，已忽略")

            print("预测参数加载成功")
            return True

        except Exception as e:
            print(f"加载YAML文件时出错: {str(e)}")
            return False

    def save_to_yaml(self, file_path: str) -> bool:
        """
        将当前参数保存到YAML文件（保留原有其他参数）

        Args:
            file_path: 保存的YAML文件路径

        Returns:
            bool: 保存成功返回True，失败返回False
        """
        try:
            # 创建目录（如果不存在）
            directory = os.path.dirname(file_path)
            if directory and not os.path.exists(directory):
                os.makedirs(directory)

            # 如果文件已存在，先加载现有配置
            existing_config = {}
            if os.path.exists(file_path):
                with open(file_path, 'r', encoding='utf-8') as file:
                    existing_config = yaml.safe_load(file) or {}

            # 合并现有配置和当前参数（当前参数覆盖现有参数）
            merged_config = existing_config.copy()
            merged_config.update(self.params)

            # 写入YAML文件
            with open(file_path, 'w', encoding='utf-8') as file:
                yaml.dump(merged_config, file, default_flow_style=False, indent=2, allow_unicode=True)

            print(f"参数已保存到: {file_path}")
            return True

        except Exception as e:
            print(f"保存YAML文件时出错: {str(e)}")
            return False


    def get_param(self, key: str) -> Any:
        """
        获取指定参数值

        Args:
            key: 参数名称

        Returns:
            参数值，如果不存在则返回None
        """
        return self.params.get(key)

    def set_param(self, key: str, value: Any) -> bool:
        """
        设置参数值

        Args:
            key: 参数名称
            value: 参数值

        Returns:
            bool: 设置成功返回True，失败返回False
        """
        if key in self.params:
            self.params[key] = value
            return True
        return False



    def get_all_params(self) -> Dict[str, Any]:
        """
        获取所有参数

        Returns:
            包含所有参数的字典
        """
        return self.params.copy()

    def reset_defaults(self):
        """
        重置为默认参数
        """
        self.params = {
            'project': 'runs/predict',
            'name': 'exp',
            'show': True,
            'save_frames': False,
            'save_txt': False,
            'save_conf': False,
            'save_crop': False,
            'show_labels': True,
            'show_conf': True,
            'show_boxes': True
        }


class BatchPredictor:
    """
    批量预测器类，用于对多个文件进行目标检测预测
    """

    def __init__(self, model_path: str = None):
        """
        初始化批量预测器

        Args:
            model_path: 模型文件路径
        """
        self.model = None
        self.model_path = model_path
        if model_path:
            self.load_model(model_path)
        self.ui_host = 'localhost'
        self.ui_port = 12346
        self.socket = None
        self.connected = False
        # 添加停止检测控制标志
        self.stop_requested = False
        self.paused = False
        self.running = True


    # 通信==============
    def connect_to_ui(self, host='localhost', port=12346):
        """连接到UI应用服务器"""
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.connect((host, port))
            self.connected = True
            print(f"已连接到UI应用 {host}:{port}")
            return True
        except Exception as e:
            print(f"连接UI应用失败: {e}")
            return False

    def send_result_to_ui(self, result_data):
        """发送检测结果到UI应用"""
        if not self.connected or not self.socket:
            print("未连接到UI应用")
            return False

        try:
            message = json.dumps(result_data, ensure_ascii=False)
            self.socket.send((message + '\n').encode('utf-8'))
            return True
        except Exception as e:
            print(f"发送数据到UI应用时出错: {e}")
            self.connected = False
            return False

    def disconnect_from_ui(self):
        """断开UI应用连接"""
        if self.socket:
            self.socket.close()
            self.connected = False

    def load_model(self, model_path: str) -> bool:
        """
        加载YOLO模型

        Args:
            model_path: 模型文件路径

        Returns:
            bool: 加载成功返回True，失败返回False
        """
        try:
            if not os.path.exists(model_path):
                print(f"错误: 模型文件不存在 - {model_path}")
                return False

            self.model = YOLO(model_path)
            self.model_path = model_path
            print(f"模型加载成功: {model_path}")
            return True

        except Exception as e:
            print(f"加载模型时出错: {str(e)}")
            return False

    def stop_detection(self):
        """
        请求停止检测过程
        """
        self.stop_requested = True
        self.running = False

        print("检测停止请求已发送")

        # 发送停止消息到UI
        stop_msg = {
            'type': 'detection_stopped',
            'message': '用户已停止检测',
            'status': 'stopped'
        }
        self.send_result_to_ui(stop_msg)


    def is_stopped(self):
        """
        检查是否请求了停止检测
        """
        return self.stop_requested

    def pause_detection(self):
        """
        暂停检测过程
        """
        self.paused = True
        print("检测已暂停")

    def resume_detection(self):
        """
        恢复检测过程
        """
        self.paused = False
        print("检测已恢复")

    def is_paused(self):
        """
        检查是否处于暂停状态
        """
        return self.paused

    def reset_detection_state(self):
        """
        重置检测状态
        """
        self.stop_requested = False
        self.paused = False

    def check_pause_stop(self):
        """
        检查暂停和停止状态，处理相应逻辑
        """
        # 检查是否请求了停止
        if self.is_stopped():
            return False  # 返回False表示需要停止
            
        # 检查是否暂停
        while self.is_paused() and not self.is_stopped():
            time.sleep(0.1)  # 短暂休眠避免占用CPU
            
        return True  # 继续执行

    def predict_batch_images(self, image_paths: List[str]):
        """
        对批量图片进行预测，并发送结果到UI应用
        """
        try:
            # 重置停止状态
            self.reset_detection_state()
            
            if not self.model:
                raise ValueError("模型未加载，请先调用load_model方法")

            # 检查所有文件是否存在
            for path in image_paths:
                if not os.path.exists(path):
                    raise FileNotFoundError(f"图片文件不存在: {path}")

            current_dir = os.path.dirname(os.path.abspath(__file__))
            # 向上两级找到项目根目录（因为当前在 UI/ui/ 下）
            project_root = os.path.dirname(os.path.dirname(current_dir))

            # 构建配置文件路径
            config_path = os.path.join(project_root, 'cfg', 'predict.yaml')
            with open(config_path, 'r') as f:
                cfg = yaml.safe_load(f)

            results_list = []
            total_images = len(image_paths)
            processed_count = 0

            # 对每张图片进行预测
            for i, image_path in enumerate(image_paths):
                # 检查暂停和停止状态
                if not self.check_pause_stop():
                    print("检测已停止")
                    stop_msg = {
                        'type': 'detection_stopped',
                        'message': '检测已被用户停止',
                        'processed': processed_count,
                        'total': total_images
                    }
                    self.send_result_to_ui(stop_msg)
                    break

                try:
                    # 执行单张图片预测
                    results = self.model(image_path, **cfg)

                    # 处理结果
                    result_data = {
                        'image_path': image_path,
                        'results': "processed",  # 简化结果表示
                        'boxes': [],
                        'classes': [],
                        'confidences': [],
                        'status': 'success',
                        # 添加统计信息
                        'total_objects': 0,
                        'class_counts': {},
                        'result_image_path': None  # 添加结果路径字段
                    }

                    # 提取检测框信息
                    if results and len(results) > 0:
                        result = results[0]  # 取第一个结果
                        if hasattr(result, 'boxes') and result.boxes is not None:
                            try:
                                boxes = result.boxes
                                result_data['boxes'] = boxes.xyxy.cpu().numpy().tolist()
                                result_data['classes'] = boxes.cls.cpu().numpy().tolist()
                                result_data['confidences'] = boxes.conf.cpu().numpy().tolist()
                            except Exception as e:
                                print(f"提取检测框信息失败: {e}")
                                result_data['boxes'] = []
                                result_data['classes'] = []
                                result_data['confidences'] = []
                            # 统计目标数量
                            result_data['total_objects'] = len(result_data['classes'])

                            # 统计各类别数量
                            class_names = getattr(self.model, 'names', {})
                            class_counts = {}
                            for cls_id in result_data['classes']:
                                cls_name = class_names.get(int(cls_id), f"Class_{int(cls_id)}")
                                class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
                            result_data['class_counts'] = class_counts

                        # 获取预测结果图片路径
                        if hasattr(result, 'path') and result.path:
                            # 从YAML配置中获取project和name参数来构造正确的保存路径
                            project_dir = cfg.get('project')
                            exp_name = cfg.get('name')

                            # 获取原始文件名
                            original_filename = os.path.basename(result.path)

                            # 构造预测结果保存路径
                            result_image_path = os.path.join(project_dir, exp_name, original_filename)
                            result_data['result_image_path'] = result_image_path

                    results_list.append(result_data)
                    processed_count += 1

                    # 发送结果到UI应用
                    self.send_result_to_ui(result_data)

                except Exception as e:
                    # 记录单张图片预测失败的情况
                    error_data = {
                        'image_path': image_path,
                        'error': str(e),
                        'status': 'failed',
                        'result_image_path': None  # 错误情况下也包含路径字段
                    }
                    results_list.append(error_data)
                    self.send_result_to_ui(error_data)

            return results_list

        except Exception as e:
            print(f"批量图片预测时出错: {str(e)}")
            error_result = {'error': str(e), 'status': 'failed'}
            self.send_result_to_ui(error_result)
            return [error_result]


    def predict_folder(self, folder_path: str):
        """
        对文件夹中的所有图片进行预测，并发送结果到UI应用
        """
        try:
            # 重置停止状态
            self.reset_detection_state()
            
            if not self.model:
                raise ValueError("模型未加载，请先调用load_model方法")

            if not os.path.exists(folder_path):
                raise FileNotFoundError(f"文件夹不存在: {folder_path}")

            # 支持的图片格式
            supported_formats = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')

            # 获取文件夹中所有图片文件
            image_paths = []
            for filename in os.listdir(folder_path):
                if filename.lower().endswith(supported_formats):
                    image_paths.append(os.path.join(folder_path, filename))

            if not image_paths:
                raise ValueError(f"文件夹中没有找到支持的图片文件: {folder_path}")

            current_dir = os.path.dirname(os.path.abspath(__file__))
            # 向上两级找到项目根目录（因为当前在 UI/ui/ 下）
            project_root = os.path.dirname(os.path.dirname(current_dir))

            # 构建配置文件路径
            config_path = os.path.join(project_root, 'cfg', 'predict.yaml')
            with open(config_path, 'r') as f:
                cfg = yaml.safe_load(f)

            # 调用批量预测方法
            return self.predict_batch_images(image_paths)

        except Exception as e:
            print(f"文件夹预测时出错: {str(e)}")
            error_result = {'error': str(e), 'status': 'failed'}
            self.send_result_to_ui(error_result)
            return [error_result]


    def predict_batch_with_progress(self, image_paths: List[str]):
        """
        带进度的批量预测，并发送结果到UI应用
        """
        try:
            # 重置停止状态
            self.reset_detection_state()

            if not self.model:
                raise ValueError("模型未加载，请先调用load_model方法")

            total_images = len(image_paths)
            if total_images == 0:
                empty_result = {'results': [], 'total': 0, 'success': 0, 'failed': 0}
                self.send_result_to_ui({'type': 'progress_summary', 'data': empty_result})
                return empty_result

            current_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(os.path.dirname(current_dir))
            config_path = os.path.join(project_root, 'cfg', 'predict.yaml')

            with open(config_path, 'r') as f:
                cfg = yaml.safe_load(f)

            results_list = []
            success_count = 0
            failed_count = 0

            # 对每张图片进行预测
            for i, image_path in enumerate(image_paths):
                # 检查暂停和停止状态
                # 在 predict_batch_with_progress 方法中找到检查暂停和停止状态的部分，修改为：
                # 检查暂停和停止状态
                if not self.check_pause_stop():
                    print("检测已停止")
                    stop_msg = {
                        'type': 'detection_stopped',
                        'message': '检测已被用户停止',
                        'processed': i,
                        'total': total_images,
                        'status': 'stopped'
                    }
                    self.send_result_to_ui(stop_msg)
                    break


                try:
                    if not os.path.exists(image_path):
                        raise FileNotFoundError(f"图片文件不存在: {image_path}")

                    # 执行单张图片预测，保存结果图片
                    results = self.model(image_path, **cfg)

                    # 处理结果
                    result_data = {
                        'image_path': image_path,
                        'results': "processed",
                        'boxes': [],
                        'classes': [],
                        'confidences': [],
                        'status': 'success',
                        'total_objects': 0,
                        'class_counts': {},
                        'result_image_path': None
                    }

                    # 提取检测框信息
                    if results and len(results) > 0:
                        result = results[0]  # 取第一个结果
                        if hasattr(result, 'boxes') and result.boxes is not None:
                            boxes = result.boxes
                            result_data['boxes'] = boxes.xyxy.cpu().numpy().tolist()  # 边界框坐标
                            result_data['classes'] = boxes.cls.cpu().numpy().tolist()  # 类别
                            result_data['confidences'] = boxes.conf.cpu().numpy().tolist()  # 置信度

                            # 统计目标数量
                            result_data['total_objects'] = len(result_data['classes'])

                            # 统计各类别数量
                            class_names = getattr(self.model, 'names', {})
                            class_counts = {}
                            for cls_id in result_data['classes']:
                                cls_name = class_names.get(int(cls_id), f"Class_{int(cls_id)}")
                                class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
                            result_data['class_counts'] = class_counts

                        # 获取预测结果图片路径
                        if hasattr(result, 'path') and result.path:
                            project_dir = cfg.get('project')
                            exp_name = cfg.get('name')
                            original_filename = os.path.basename(result.path)
                            result_image_path = os.path.join(project_dir, exp_name, original_filename)
                            result_data['result_image_path'] = result_image_path
                            print("结果路径: " + result_image_path)

                    results_list.append(result_data)
                    success_count += 1

                    # 发送单个结果到UI应用
                    self.send_result_to_ui(result_data)

                except Exception as e:
                    # 记录单张图片预测失败的情况
                    error_data = {
                        'image_path': image_path,
                        'error': str(e),
                        'status': 'failed'
                    }
                    results_list.append(error_data)
                    failed_count += 1
                    self.send_result_to_ui(error_data)

                # 发送进度信息到UI应用
                progress_data = {
                    'type': 'progress',
                    'current': i + 1,
                    'total': total_images,
                    'percentage': (i + 1) / total_images * 100,
                    'success': success_count,
                    'failed': failed_count
                }
                self.send_result_to_ui(progress_data)

            # 只有在没有被停止的情况下才发送总结信息
            if not self.is_stopped():
                summary_data = {
                    'type': 'progress_summary',
                    'results': results_list,
                    'total': total_images,
                    'success': success_count,
                    'failed': failed_count,
                    'progress': 100.0
                }
            else:
                summary_data = {
                    'type': 'progress_summary',
                    'results': results_list,
                    'total': total_images,
                    'success': success_count,
                    'failed': failed_count,
                    'progress': (success_count + failed_count) / total_images * 100 if total_images > 0 else 0,
                    'stopped': True
                }

            # 发送总结信息到UI应用
            self.send_result_to_ui(summary_data)
            return summary_data

        except Exception as e:
            print(f"带进度的批量预测时出错: {str(e)}")
            error_result = {
                'type': 'progress_error',
                'results': [{'error': str(e), 'status': 'failed'}],
                'total': 0,
                'success': 0,
                'failed': 1,
                'progress': 0.0
            }
            self.send_result_to_ui(error_result)
            return error_result


    def predict_single_image(self, image_path: str):
        """
        对单张图片进行预测，并发送结果到UI应用
        """
        try:
            # 重置停止状态
            self.reset_detection_state()
            
            if not self.model:
                raise ValueError("模型未加载，请先调用load_model方法")

            if not os.path.exists(image_path):
                raise FileNotFoundError(f"图片文件不存在: {image_path}")

            current_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(os.path.dirname(current_dir))
            config_path = os.path.join(project_root, 'cfg', 'predict.yaml')

            with open(config_path, 'r') as f:
                cfg = yaml.safe_load(f)

            # 确保保存结果
            cfg['save'] = True

            # 检查暂停和停止状态
            if not self.check_pause_stop():
                stop_msg = {
                    'type': 'detection_stopped',
                    'message': '检测已被用户停止',
                    'processed': 0,
                    'total': 1,
                    'status': 'stopped'
                }
                self.send_result_to_ui(stop_msg)
                return

            # 执行单张图片预测
            results = self.model(image_path, **cfg)

            # 处理结果
            result_data = {
                'image_path': image_path,
                'results': "processed",
                'boxes': [],
                'classes': [],
                'confidences': [],
                'status': 'success',
                'total_objects': 0,
                'class_counts': {},
                'result_image_path': None
            }

            # 提取检测框信息
            if results and len(results) > 0:
                result = results[0]  # 取第一个结果
                if hasattr(result, 'boxes') and result.boxes is not None:
                    boxes = result.boxes
                    result_data['boxes'] = boxes.xyxy.cpu().numpy().tolist()  # 边界框坐标
                    result_data['classes'] = boxes.cls.cpu().numpy().tolist()  # 类别
                    result_data['confidences'] = boxes.conf.cpu().numpy().tolist()  # 置信度

                    # 统计目标数量
                    result_data['total_objects'] = len(result_data['classes'])

                    # 统计各类别数量
                    class_names = getattr(self.model, 'names', {})
                    class_counts = {}
                    for cls_id in result_data['classes']:
                        cls_name = class_names.get(int(cls_id), f"Class_{int(cls_id)}")
                        class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
                    result_data['class_counts'] = class_counts

                # 获取预测结果图片路径
                if hasattr(result, 'path') and result.path:
                    # 从YAML配置中获取project和name参数来构造正确的保存路径
                    project_dir = cfg.get('project')
                    exp_name = cfg.get('name')

                    # 获取原始文件名
                    original_filename = os.path.basename(result.path)

                    # 构造预测结果保存路径
                    result_image_path = os.path.join(project_dir, exp_name, original_filename)
                    result_data['result_image_path'] = result_image_path
                    print("结果路径: " + result_image_path)

                else:
                    # 如果没有返回路径，尝试构造默认路径
                    # YOLO默认保存在 runs/detect/exp/ 目录下
                    filename = os.path.basename(image_path)
                    default_save_dir = os.path.join(os.path.dirname(image_path), 'runs', 'detect', 'exp')
                    default_result_path = os.path.join(default_save_dir, filename)
                    if os.path.exists(default_result_path):
                        result_data['result_image_path'] = default_result_path

            # 发送结果到UI应用
            self.send_result_to_ui(result_data)
            return result_data

        except Exception as e:
            print(f"单张图片预测时出错: {str(e)}")
            error_result = {'error': str(e), 'status': 'failed'}
            self.send_result_to_ui(error_result)
            return error_result


    def predict_video(self, video_path: str):
        try:
            # 重置停止状态
            self.reset_detection_state()

            if not self.model:
                raise ValueError("模型未加载，请先调用load_model方法")

            if not os.path.exists(video_path):
                raise FileNotFoundError(f"视频文件不存在: {video_path}")

            current_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(os.path.dirname(current_dir))
            config_path = os.path.join(project_root, 'cfg', 'predict.yaml')

            with open(config_path, 'r') as f:
                cfg = yaml.safe_load(f)

            # 设置视频预测参数
            cfg['save'] = True
            cfg['stream'] = True

            # 添加跳帧参数，例如每5帧处理一帧（减少80%的帧数）
            frame_skip = 5  # 可以根据需要调整这个值，1表示不跳帧，2表示跳50%的帧，5表示跳80%的帧

            # 发送开始视频预测消息
            start_msg = {
                'type': 'video_start',
                'video_path': video_path,
                'status': 'processing'
            }
            self.send_result_to_ui(start_msg)

            # 执行视频预测
            results = self.model(video_path, **cfg)

            frame_count = 0
            processed_frame_count = 0
            success_count = 0
            failed_count = 0

            # 处理每一帧
            for result in results:
                # 检查暂停和停止状态 - 添加更频繁的检查
                if not self.check_pause_stop():
                    print("视频检测已停止")
                    stop_msg = {
                        'type': 'detection_stopped',
                        'message': '视频检测已被用户停止',
                        'frame_count': frame_count,
                        'status': 'stopped'
                    }
                    self.send_result_to_ui(stop_msg)
                    break

                try:
                    frame_count += 1

                    # 跳帧处理：只处理每隔frame_skip帧
                    if (frame_count - 1) % frame_skip != 0:
                        continue  # 跳过这一帧

                    processed_frame_count += 1

                    # 发送帧处理进度 - 减少发送频率，每10帧发送一次
                    if processed_frame_count % 10 == 0:
                        progress_msg = {
                            'type': 'video_progress',
                            'processed_frames': processed_frame_count,
                            'total_frames': frame_count,
                            'percentage': (processed_frame_count / frame_count * 100) if frame_count > 0 else 0
                        }
                        self.send_result_to_ui(progress_msg)

                    # 处理帧数据
                    frame_data = {
                        'type': 'video_frame',
                        'frame_number': processed_frame_count,
                        'original_frame_number': frame_count,  # 添加原始帧号信息
                        'status': 'success'
                    }

                    # 提取检测框信息
                    if hasattr(result, 'boxes') and result.boxes is not None:
                        boxes = result.boxes
                        frame_data['boxes'] = boxes.xyxy.cpu().numpy().tolist()
                        frame_data['classes'] = boxes.cls.cpu().numpy().tolist()
                        frame_data['confidences'] = boxes.conf.cpu().numpy().tolist()

                        # 统计目标数量
                        frame_data['total_objects'] = len(frame_data['classes'])

                        # 统计各类别数量
                        class_names = getattr(self.model, 'names', {})
                        class_counts = {}
                        for cls_id in frame_data['classes']:
                            cls_name = class_names.get(int(cls_id), f"Class_{int(cls_id)}")
                            class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
                        frame_data['class_counts'] = class_counts

                    # 获取帧图片路径
                    if hasattr(result, 'path') and result.path:
                        frame_data['result_image_path'] = result.path

                    # 减少发送到UI的频率，只发送每5帧的结果
                    if processed_frame_count % 5 == 0:
                        self.send_result_to_ui(frame_data)

                    success_count += 1

                except Exception as frame_e:
                    failed_count += 1
                    error_frame_data = {
                        'type': 'video_frame',
                        'frame_number': processed_frame_count,
                        'original_frame_number': frame_count,
                        'error': str(frame_e),
                        'status': 'failed'
                    }
                    self.send_result_to_ui(error_frame_data)
                    print(f"处理视频帧 {processed_frame_count} (原始帧 {frame_count}) 时出错: {str(frame_e)}")

            # 发送视频处理完成消息
            # 构造预测结果视频路径 - 修复路径问题
            project_dir = cfg.get('project')
            exp_name = cfg.get('name')
            video_filename = os.path.basename(video_path)

            # 保持原始文件扩展名
            result_video_path = os.path.join(project_dir, exp_name, video_filename)

            # 如果project_dir是相对路径，转换为绝对路径
            if not os.path.isabs(project_dir):
                result_video_path = os.path.join(project_root, result_video_path)

            # 如果结果文件不存在，尝试其他可能的扩展名和路径
            if not os.path.exists(result_video_path):
                # 尝试不同的扩展名
                base_name = os.path.splitext(video_filename)[0]
                dir_path = os.path.dirname(result_video_path)
                possible_extensions = ['.mp4', '.avi', '.mov', '.mkv', '.wmv']

                # 尝试当前构造的路径
                found = False
                for ext in possible_extensions:
                    test_path = os.path.join(dir_path, base_name + ext)
                    if os.path.exists(test_path):
                        result_video_path = test_path
                        found = True
                        break

                # 如果还没找到，尝试其他可能的路径
                if not found:
                    # 尝试在UI/ui子目录中查找
                    ui_ui_dir = os.path.join(project_root, 'UI', 'ui')
                    ui_project_dir = os.path.join(ui_ui_dir, project_dir) if not os.path.isabs(project_dir) else project_dir
                    ui_result_video_path = os.path.join(ui_project_dir, exp_name, video_filename)

                    if os.path.exists(ui_result_video_path):
                        result_video_path = ui_result_video_path
                    else:
                        # 尝试不同的扩展名在UI/ui路径中
                        ui_dir_path = os.path.dirname(ui_result_video_path)
                        for ext in possible_extensions:
                            test_path = os.path.join(ui_dir_path, base_name + ext)
                            if os.path.exists(test_path):
                                result_video_path = test_path
                                break

            # 发送视频结束消息
            end_msg = {
                'type': 'video_end',
                'video_path': video_path,
                'result_video_path': result_video_path,
                'total_frames': frame_count,
                'processed_frames': processed_frame_count,  # 添加实际处理的帧数
                'success_frames': success_count,
                'failed_frames': failed_count,
                'frame_skip': frame_skip,  # 添加跳帧信息
                'status': 'completed'
            }
            self.send_result_to_ui(end_msg)

            # 发送视频完成消息，触发自动播放
            finished_msg = {
                'type': 'video_finished',
                'video_path': result_video_path,  # 使用修复后的路径
                'status': 'completed'
            }
            self.send_result_to_ui(finished_msg)

            return end_msg

        except Exception as e:
            print(f"视频预测时出错: {str(e)}")
            error_result = {
                'type': 'video_error',
                'video_path': video_path,
                'error': str(e),
                'status': 'failed'
            }
            self.send_result_to_ui(error_result)
            return error_result

    def predict_camera(self):
        """
        对摄像头进行实时预测，并发送结果到UI应用
        """
        try:
            # 重置停止状态
            self.reset_detection_state()

            if not self.model:
                raise ValueError("模型未加载，请先调用load_model方法")

            current_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(os.path.dirname(current_dir))
            config_path = os.path.join(project_root, 'cfg', 'predict.yaml')

            with open(config_path, 'r') as f:
                cfg = yaml.safe_load(f)

            # 设置摄像头预测参数
            cfg['save'] = True  # 改为True以保存检测结果
            cfg['stream'] = True
            cfg['show'] = False  # 不显示窗口

            # 添加跳帧参数，例如每3帧处理一帧
            frame_skip = 3  # 可根据需要调整
            frame_count = 0

            # 发送开始摄像头预测消息
            start_msg = {
                'type': 'camera_start',
                'status': 'processing'
            }
            self.send_result_to_ui(start_msg)

            # 执行摄像头预测 (source=0 表示默认摄像头)
            results = self.model(source=0, **cfg)

            processed_frame_count = 0

            # 将处理帧的逻辑改为异步处理
            def process_frames():
                nonlocal frame_count, processed_frame_count
                try:
                    # 获取下一个结果（在流模式下这应该是非阻塞的）
                    try:
                        result = next(results)
                    except StopIteration:
                        # 摄像头流结束
                        return

                    # 检查暂停和停止状态
                    if not self.check_pause_stop():
                        print("摄像头检测已停止")
                        stop_msg = {
                            'type': 'camera_stopped',
                            'message': '摄像头检测已被用户停止',
                            'frame_count': processed_frame_count
                        }
                        self.send_result_to_ui(stop_msg)
                        return

                    try:
                        frame_count += 1

                        # 跳帧处理：只处理每隔frame_skip帧
                        if (frame_count - 1) % frame_skip != 0:
                            # 继续处理下一帧
                            if not self.is_stopped():
                                # 使用after_idle让出控制权，然后继续处理下一帧
                                import threading
                                timer = threading.Timer(0.01, process_frames)
                                timer.daemon = True
                                timer.start()
                            return  # 跳过这一帧

                        processed_frame_count += 1

                        # 发送帧处理进度
                        frame_data = {
                            'type': 'camera_frame',
                            'frame_number': processed_frame_count,
                            'original_frame_number': frame_count,  # 添加原始帧号
                            'status': 'success'
                        }

                        # 提取检测框信息
                        if hasattr(result, 'boxes') and result.boxes is not None:
                            boxes = result.boxes
                            frame_data['boxes'] = boxes.xyxy.cpu().numpy().tolist()
                            frame_data['classes'] = boxes.cls.cpu().numpy().tolist()
                            frame_data['confidences'] = boxes.conf.cpu().numpy().tolist()

                            # 统计目标数量
                            frame_data['total_objects'] = len(frame_data['classes'])

                            # 统计各类别数量
                            class_names = getattr(self.model, 'names', {})
                            class_counts = {}
                            for cls_id in frame_data['classes']:
                                cls_name = class_names.get(int(cls_id), f"Class_{int(cls_id)}")
                                class_counts[cls_name] = class_counts.get(cls_name, 0) + 1
                            frame_data['class_counts'] = class_counts

                        # 获取检测结果图片路径
                        if hasattr(result, 'path') and result.path:
                            # 从YAML配置中获取project和name参数来构造正确的保存路径
                            project_dir = cfg.get('project')
                            exp_name = cfg.get('name')

                            # 获取原始文件名
                            original_filename = os.path.basename(result.path)

                            # 构造预测结果保存路径
                            result_image_path = os.path.join(project_dir, exp_name, original_filename)
                            frame_data['result_image_path'] = result_image_path
                            print("摄像头检测结果路径: " + result_image_path)

                        self.send_result_to_ui(frame_data)

                    except Exception as frame_e:
                        error_frame_data = {
                            'type': 'camera_frame',
                            'frame_number': processed_frame_count,
                            'original_frame_number': frame_count,
                            'error': str(frame_e),
                            'status': 'failed'
                        }
                        self.send_result_to_ui(error_frame_data)

                    # 继续处理下一帧（如果未停止）
                    if not self.is_stopped():
                        # 使用线程定时器让出控制权，然后继续处理下一帧
                        import threading
                        timer = threading.Timer(0.01, process_frames)
                        timer.daemon = True
                        timer.start()

                except Exception as e:
                    print(f"处理摄像头帧时出错: {str(e)}")
                    error_result = {
                        'type': 'camera_error',
                        'error': str(e),
                        'status': 'failed'
                    }
                    self.send_result_to_ui(error_result)

            # 开始异步处理帧
            process_frames()

        except Exception as e:
            print(f"摄像头预测时出错: {str(e)}")
            error_result = {
                'type': 'camera_error',
                'error': str(e),
                'status': 'failed'
            }
            self.send_result_to_ui(error_result)
            return error_result





    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息，并发送到UI应用
        """
        if not self.model:
            error_info = {'error': '模型未加载'}
            self.send_result_to_ui({'type': 'model_info', 'data': error_info})
            return error_info

        try:
            info = {
                'model_path': self.model_path,
                'task': getattr(self.model, 'task', 'unknown'),
                'names': getattr(self.model, 'names', {})
            }
            self.send_result_to_ui({'type': 'model_info', 'data': info})
            return info
        except Exception as e:
            error_info = {'error': str(e)}
            self.send_result_to_ui({'type': 'model_info', 'data': error_info})
            return error_info
